#!/usr/bin/python
from glob import glob
from SocketServer import UnixStreamServer, StreamRequestHandler
from StringIO import StringIO
from struct import unpack, pack, calcsize
from subprocess import Popen
from threading import Thread, Lock
from time import time, sleep
import errno
import logging
import os
import pickle
import socket
import sys
import signal

# Crabs are known for their remote process calls
LENGTH_STRUCT_FMT = "Q"
LENGTH_STRUCT_LENGTH = calcsize(LENGTH_STRUCT_FMT)

if __name__ != "__main__":
    # You don't have access to the config if you are the child process
    from config import config
    REMOTE_FILE_HANDLER_TIMEOUT = config.getint("irs", "remote_file_handler_timeout")
else:
    sys.path.append(os.path.join(os.path.dirname(sys.argv[0]), "../"))
    import storage.misc as misc

import constants

class CrabRPCServer(object):
    def __init__(self, path):
        self._handler = CrabHandler
        self._server = UnixStreamServer(path, self._handler)
        self.registeredFunctions = {}

    def register_function(self, func, name=None):
        if name is None:
            name = func.__name__

        if not hasattr(self._handler, "registeredFunctions"):
            self._handler.registeredFunctions = {}
        self._handler.registeredFunctions[name] = func

    def serve_forever(self):
        self._server.serve_forever()

class CrabHandler(StreamRequestHandler):
    _log = logging.getLogger("CrabHandler")
    def handle(self):
        while True:
            try:
                rawLength = self.rfile.read(LENGTH_STRUCT_LENGTH)
                length = unpack(LENGTH_STRUCT_FMT, rawLength)[0]
                pickledCall = self.rfile.read(length)
                name, args, kwargs = pickle.loads(pickledCall)
                err = res = None
                try:
                    res = self.callRegisteredFunction(name, args, kwargs)
                except Exception, ex:
                    err = ex

                resp = pickle.dumps((res, err))
                self.wfile.write(pack(LENGTH_STRUCT_FMT, len(resp)))
                self.wfile.flush()
                self.wfile.write(resp)
                self.wfile.flush()
            except:
                self._log.warn("Could not complete operation", exc_info=True)

    def callRegisteredFunction(self, name, args, kwargs):
        return self.registeredFunctions[name](*args, **kwargs)

class CrabRPCProxy(object):
    def __init__(self, path, timeoutCallback):
        self._timeoutCallback = timeoutCallback
        self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self._sockPath = path
        lastError = Exception("Something odd happened")
        for i in range(5):
            try:
                self._sock.connect(self._sockPath)
                return
            except Exception, ex:
                sleep(1)
                lastError = ex
        raise lastError

    def _recvAll(self, length, timeout):
        self._sock.settimeout(timeout)
        startTime = time()
        rawResponse = StringIO()
        while rawResponse.tell() < length:
            if (time() - startTime) > timeout:
                raise socket.timeout()

            rawResponse.write(self._sock.recv(length - rawResponse.tell()))

        rawResponse.seek(0)
        return rawResponse.read()

    def _sendAll(self, data):
        self._sock.settimeout(REMOTE_FILE_HANDLER_TIMEOUT)
        l = 0
        startTime = time()
        while l < len(data):
            if (time() - startTime) > REMOTE_FILE_HANDLER_TIMEOUT:
                raise socket.timeout()
            l += self._sock.send(data[l:])

    def __getattr__(self, name):
        def callCrabRPCFunction (*args, **kwargs):
            request = pickle.dumps((name, args, kwargs))
            self._sendAll(pack(LENGTH_STRUCT_FMT, len(request)))
            self._sendAll(request)

            try:
                rawLength = self._recvAll(LENGTH_STRUCT_LENGTH, REMOTE_FILE_HANDLER_TIMEOUT)
            except socket.timeout:
                self._timeoutCallback()
                rawLength = self._recvAll(LENGTH_STRUCT_LENGTH, 2)
            length = unpack(LENGTH_STRUCT_FMT, rawLength)[0]
            rawResponse = self._recvAll(length, REMOTE_FILE_HANDLER_TIMEOUT)


            res, err = pickle.loads(rawResponse)
            if err is not None:
                raise err
            return res

        return callCrabRPCFunction

class PoolHandler (object):
    def __init__(self, index):
        address = "/var/vdsm/rfh%d.sock" % (index)
        if os.path.exists(address):
            os.unlink(address)
        self.process = Popen([constants.EXT_PYTHON, __file__, str(os.getpid()), address], close_fds=True)
        self.proxy = CrabRPCProxy(address, lambda : os.kill(self.process.pid, signal.SIGINT))

class _RemoteFileHandlerPool(object):
    _log = logging.getLogger("Storage.RemoteFileHandler")
    def __init__(self, numOfHandlers):
        self._numOfHandlers = numOfHandlers
        self.handlers = [None] * numOfHandlers
        self.occupied = [Lock() for i in xrange(numOfHandlers)]

    def _isHandlerAvailable(self, poolHandler):
        if poolHandler is None:
            return False
        try:
            os.kill(poolHandler.process.pid, 0)
        except OSError:
            return False
        return True

    def __getattr__(self, name):
        def wrapper(*args, **kwargs):
            for i, isOccupied in enumerate(self.occupied):
                if not isOccupied.acquire(False):
                    continue

                try:
                    handler = self.handlers[i]
                    if not self._isHandlerAvailable(handler):
                        handler = self.handlers[i] = PoolHandler(i)

                    try:
                        return getattr(handler.proxy, name)(*args, **kwargs)
                    except socket.timeout:
                        try:
                            os.kill(handler.process.pid, signal.SIGKILL)
                        except:
                            self._log("Could not signal stuck handler (PID:%d)", handler.process.pid, exc_info=True)

                        self.handlers[i] = None
                        err = OSError()
                        err.errno = errno.ETIME
                        err.strerror = "Operation stuck on remote handler"
                        raise err
                    except KeyboardInterrupt:
                        err = OSError()
                        err.errno = errno.ETIME
                        err.strerror = "Operation was stuck on remote handler but I snapped it out of it"
                        raise err
                finally:
                    isOccupied.release()

            raise Exception("No free file handlers in pool")

        return wrapper


def pokeParent():
    try:
        while True:
            os.kill(parentPid, 0)
            sleep(2)
    except:
        os.kill(os.getpid(), signal.SIGKILL)

def writeLines(path, data):
    f = open(path, "w")
    try:
        f.writelines(data)
    finally:
        f.close()

def write(path, data):
    f = open(path, "w")
    try:
        f.write(data)
    finally:
        f.close()

def readLines(path):
    f = open(path, "r")
    try:
        return f.readlines()
    finally:
        f.close()

def umount(mountPoint, force=False):
    options = []
    if force:
        options.append("-f")

    umount_cmd = [constants.EXT_UMOUNT] + options + [mountPoint]
    (rc, out, err) = misc.execCmd(umount_cmd)
    return rc

def touch(path):
    f = open(path, "w")
    f.close()

if __name__ == "__main__":
    #Parse args
    try:
        parentPid, sockAddr = sys.argv[1:]
        parentPid = int(parentPid)
    except ValueError, ex:
        sys.stderr.write("Error parsing args %s\n" % ex)
        sys.exit(errno.EINVAL)

    try:
        try:
            os.unlink(sockAddr)
        except OSError:
            pass
        server = CrabRPCServer(sockAddr)
        server.register_function(glob)
        server.register_function(os.access)
        server.register_function(os.chmod)
        server.register_function(os.link)
        server.register_function(os.listdir)
        server.register_function(os.makedirs)
        server.register_function(os.mkdir)
        server.register_function(os.path.abspath)
        server.register_function(os.path.exists)
        server.register_function(os.path.isdir)
        server.register_function(os.path.ismount)
        server.register_function(os.path.lexists)
        server.register_function(os.rename)
        server.register_function(os.stat)
        server.register_function(os.statvfs)
        server.register_function(os.symlink)
        server.register_function(os.unlink)
        server.register_function(readLines)
        server.register_function(touch)
        server.register_function(umount)
        server.register_function(write)
        server.register_function(writeLines)
    except Exception, ex:
        sys.stderr.write("Error creating CrabRPC server %s\n" % ex)
        sys.exit(errno.ENOEXEC)

    pokeThread = Thread(target=pokeParent)
    pokeThread.setDaemon(True)
    pokeThread.start()

    try:
        server.serve_forever()
    except Exception, ex:
        sys.stderr.write("Error while serving %s\n" % ex)
        sys.exit(errno.ENOEXEC)
else:
    remoteFileHandlerPool = _RemoteFileHandlerPool(config.getint("irs", "remote_file_handlers_num"))
